import casadi as ca

class IrisEOM:
    def __init__(self):
        """
        初始化IRIS无人机系统模型参数
        """
        self.J = 0.029125  # inertia
        self.B = 0.4       # arm_length
        self.tau = 0.0125  # time_constant

        self.rotor_constant = 5.84e-06
        self.hover_pwm = 500
        
        # 创建系统的CasADi函数
        self._create_system_model()
        
    def _create_system_model(self):
        """
        创建系统的数学模型
        状态方程: 
        dy(1) = y(2)
        dy(2) = -(1/tau)*y(2) + (B/(J*tau))*u
        """
        # 定义符号变量
        y1 = ca.MX.sym('y1')  # 位置
        y2 = ca.MX.sym('y2')  # 速度
        u = ca.MX.sym('u')    # 控制输入
        
        # 状态向量
        y = ca.vertcat(y1, y2)

        dy1 = y2
        dy2 = -(1/self.tau)*y2 + (self.B/(self.J*self.tau))*u
        dy = ca.vertcat(dy1, dy2)

        self.iris_eom = ca.Function(
            'iris_eom',
            [y, u],
            [dy],
            ['state', 'input'],
            ['dstate']
        )
    
    def set_rotor_input(self, u):
        thrust = self.rotor_constant * self.hover_pwm**2 * u * abs(u)
        return thrust
        
    
